{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Recognition using TAO Metric Learning Recognition\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a model for object recogtion on an [ImageNet](https://www.image-net.org/) format classification dataset.\n",
    "* Evaluate the trained model & export results.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to an .onnx file for deployment to DeepStream or TensorRT.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained `MLRecog` model which you may deploy via [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of MLRecogNet using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Inferences](#head-6)\n",
    "7. [Deploy](#head-7)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"ml_recognition\", \"results\")\n",
    "os.environ[\"HOST_MODEL_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"ml_recognition\", \"models\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/metric_learning_recognition\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR\n",
    "! mkdir -p $HOST_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tao_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_MODEL_DIR\"],\n",
    "           \"destination\": \"/model\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tao_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "Here we use [Retail Product Checkout dataset](https://www.kaggle.com/datasets/diyer22/retail-product-checkout-dataset) to illustrate the method of training the metric recognition model for retail item recognition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [Action required] Download the dataset manually.\n",
    "# [Action required] Put your downloaded .zip dataset file at $HOST_DATA_DIR/retail-product-checkout-dataset.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Extract the files\n",
    "# apt-get install unzip\n",
    "!mkdir $HOST_DATA_DIR\n",
    "\n",
    "# set dataset root folder path\n",
    "%env DATA_FOLDER=retail-product-checkout-dataset_classification_demo\n",
    "\n",
    "# Run data processing script: \n",
    "# 1. crop the images and save as a classification dataset\n",
    "# 2. split the dataset as train/val/test/reference sets\n",
    "# 3. separate the classes to be known and unknown classes\n",
    "\n",
    "# install the pkgs needed for process script if needed\n",
    "!pip install Cython==0.29.36\n",
    "!pip install opencv-python\n",
    "!pip install pycocotools\n",
    "!pip install tqdm\n",
    "# now run the process script\n",
    "!python $NOTEBOOK_ROOT/process_retail_product_checkout_dataset.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Verify\n",
    "!ls -l $HOST_DATA_DIR/$DATA_FOLDER/known_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $HOST_DATA_DIR/$DATA_FOLDER/unknown_classes"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    " \n",
    "* results_dir: a global setup for output directories, would create train/evaluation/inference/export subdirectories based on subtasks. Can be overwritten by subtask ``results_dir`` fields.\n",
    "\n",
    "* model: configure the model setting\n",
    "  * backbone: type of backbone architecture, supported backbones: `resnet_50`, `resnet_101`, `fan_tiny`, `fan_small`, `fan_base`, `fan_large`, `nvdinov2_vit_large_legacy`\n",
    "  * pretrained_model_path: path for the pretrained model weights\n",
    "  * pretrained_trunk_path: path for trunk pretrained weights\n",
    "  * pretrained_embedder_path: path for embedder pretrained weights\n",
    "  * input_width: width of an input image\n",
    "  * input_height: height of an input image\n",
    "  * input_channels: number of color channels for input images, always in channel first format\n",
    "  * feat_dim: size of the output embedding\n",
    "\n",
    "* train: configure the training hyperparameters\n",
    "  * train_trunk: If false, the trunk parameters will be frozen. Default true.\n",
    "  * train_embedder: If false, the embedder parameters will be frozen. Default true.\n",
    "  * optim: configure optimizer\n",
    "  * num_epochs: number of epochs\n",
    "  * checkpoint_interval: enabling how often to store models\n",
    "  * grad_clip: enabling gradient clipping\n",
    "  * smooth_loss: enabling label smoothing feature, True/False\n",
    "  * batch_size: number of images in 1 batch for training\n",
    "  * val_batch_size: number of images in 1 batch for validation\n",
    "  * resume_training_checkpoint_path: resume .pth model training from a saved checkpoint\n",
    "  * report_accuracy_per_class: enabling accuracy per class report instead of average class accuracies, True/False\n",
    "  \n",
    "* dataset: configure the dataset and augmentation methods\n",
    "  * train_dataset: path for the train dataset directory\n",
    "  * val_dataset: map of the validation or test dataset directory. It contains reference and query set.\n",
    "  * workers: number of workers to do data loading\n",
    "  * pixel_mean: pixel mean in 3 channels for normalization\n",
    "  * pixel_std: pixel standard deviation in 3 channels for normalization\n",
    "  * prob: probability of randomly flipping images horizontally\n",
    "  * re_prob: constant for random erasing\n",
    "  * gaussian_blur: configurations for gaussian blur\n",
    "  * color_augmentation: configurations for color augmentation\n",
    "  * num_instance: number of types 1 image is repeated in a batch\n",
    "  * class_map: path to the yaml file mapping dataset class name to the new class names\n",
    "\n",
    "* evaluate: configure evaluate subtask parameters\n",
    "  * checkpoint: the .pth model for evaluation\n",
    "  * trt_engine: path of the tensorrt engine for evaluate\n",
    "  * report_accuracy_per_class: enabling accuracy per class report instead of average class accuracies, True/False\n",
    "  * topk: get predictions by the k nearest neighbor\n",
    "  * batch_size: the batch size for evaluate\n",
    "  * results_dir: the evaluation output directory. Have priority over global `results_dir`\n",
    "\n",
    "* inference: configure inference subtask parameters\n",
    "  * inference_input_type: the format of query dataset, image/image_folder/classification_folder\n",
    "  * checkpoint: the .pth model for inference\n",
    "  * trt_engine: path of the tensorrt engine for inference\n",
    "  * input_path: the inference image/image folder/classification dataset folder\n",
    "  * topk: get predictions by the k nearest neighbors\n",
    "  * batch_size: the batch size for inference\n",
    "  * results_dir: the inference output directory. Have priority over global `results_dir`\n",
    " \n",
    "* export: configure export subtask parameters\n",
    "  * checkpoint: the .pth model for export (to onnx file)\n",
    "  * onnx_file: the exported onnx model path. Have priority over the default onnx name created from ``export.results_dir``. \n",
    "  * gpu_id: the index of a single GPU only for export. default 0.\n",
    " \n",
    " \n",
    "* gen_trt_engine: configure tensorrt generation subtask parameters\n",
    "  * gpu_id: the index of a single GPU only for tensorrt engine generation. default 0.\n",
    "  * onnx_file: path of the onnx file that tensorrt engine converted from\n",
    "  * trt_engine: path of the tensorrt engine to generate\n",
    "  * batch_size: the batch size of the tensorrt engine. When `batch_size=-1`, a dynamic batch size trt engine would be generated.\n",
    "  * verbose: If True, verbose information of tensorrt generation would be printed out\n",
    "  * tensorrt: trt engine generation setup\n",
    "  * results_dir: the trt engine generation output folder. Have priority over global `results_dir`\n",
    "\n",
    "Please refer to the [TAO documentation - Metric Learning Recognition](https://docs.nvidia.com/tao/tao-toolkit/text/metric_learning_recognition/metric_learning_recognition.html) to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models.\n",
    "* WARNING: Training will take several hours or one day to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Train Metric Learning Recognition model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will train a MLRecog model with ResNet101 backbone and 2048 embedding size output. The backbone would be loaded with weights trained by NVImageNetV2 (same classes as [ImageNet](https://www.image-net.org/) but using licensed datasets)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env EPOCH=149"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train model\")\n",
    "! tao model ml_recog train \\\n",
    "              -e $SPECS_DIR/train.yaml \\\n",
    "              results_dir=$RESULTS_DIR \\\n",
    "              dataset.train_dataset=$DATA_DIR/$DATA_FOLDER/known_classes/train \\\n",
    "              dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "              dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## Training command for multi-gpu training. We can define the number of gpus and specify which GPU's are to be used by setting the `train.gpu_ids` parameter.\n",
    "## The following command will trigger multi-gpu training on gpu 0 and gpu 1.\n",
    "# ! tao model ml_recog train \\\n",
    "#               -e $SPECS_DIR/train.yaml \\\n",
    "#               results_dir=$RESULTS_DIR \\\n",
    "#               dataset.train_dataset=$DATA_DIR/$DATA_FOLDER/known_classes/train \\\n",
    "#               dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "#               dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/val \\\n",
    "#               train.gpu_ids=[0,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/ml_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/resnet101_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/resnet101_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>\n",
    "Evaluate trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate on known classes\n",
    "! tao model ml_recog evaluate \\\n",
    "            -e $SPECS_DIR/evaluate.yaml \\\n",
    "            evaluate.results_dir=$RESULTS_DIR/evaluate \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/resnet101_model.pth \\\n",
    "            dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "            dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/test \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate on unknown classes\n",
    "! tao model ml_recog evaluate \\\n",
    "            -e $SPECS_DIR/evaluate.yaml \\\n",
    "            evaluate.results_dir=$RESULTS_DIR/eval_unknown \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/resnet101_model.pth \\\n",
    "            dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/unknown_classes/reference \\\n",
    "            dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/unknown_classes/test \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Inferences <a class=\"anchor\" id=\"head-6\"></a>\n",
    "In this section, we run the metric_learning inference tool to generate inferences with the trained models and save the results under `$RESULTS_DIR`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# run inference on known classes\n",
    "! tao model ml_recog inference \\\n",
    "                    -e $SPECS_DIR/infer.yaml \\\n",
    "                    inference.results_dir=$RESULTS_DIR/inference \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/resnet101_model.pth \\\n",
    "                    dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "                    inference.input_path=$DATA_DIR/$DATA_FOLDER/known_classes/test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run inference on unknown classes\n",
    "! tao model ml_recog inference \\\n",
    "                    -e $SPECS_DIR/infer.yaml \\\n",
    "                    inference.results_dir=$RESULTS_DIR/inference_unknown \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/resnet101_model.pth \\\n",
    "                    dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/unknown_classes/reference \\\n",
    "                    inference.input_path=$DATA_DIR/$DATA_FOLDER/unknown_classes/test "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Inference results:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/inference\n",
    "!ls -lth $HOST_RESULTS_DIR/inference_unknown"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Deploy <a class=\"anchor\" id=\"head-7\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Export the model to ONNX model.\n",
    "! tao model ml_recog export \\\n",
    "                   -e $SPECS_DIR/export.yaml \\\n",
    "                   export.results_dir=$RESULTS_DIR/export \\\n",
    "                   export.checkpoint=$RESULTS_DIR/train/resnet101_model.pth \\\n",
    "                   export.onnx_file=$RESULTS_DIR/export/resnet101_model.onnx\n",
    "               "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate TensorRT engine using tao deploy\n",
    "!tao deploy ml_recog gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \\\n",
    "                                       gen_trt_engine.onnx_file=$RESULTS_DIR/export/resnet101_model.onnx \\\n",
    "                                       gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/resnet101_model.engine \\\n",
    "                                       results_dir=$RESULTS_DIR\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate int8 TensorRT engine using tao deploy\n",
    "!tao deploy ml_recog gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \\\n",
    "                                       gen_trt_engine.onnx_file=$RESULTS_DIR/export/resnet101_model.onnx \\\n",
    "                                       gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/resnet101_model.int8.engine \\\n",
    "                                       results_dir=$RESULTS_DIR \\\n",
    "                                       gen_trt_engine.tensorrt.data_type=int8 \\\n",
    "                                       gen_trt_engine.tensorrt.calibration.cal_image_dir=[$DATA_DIR/$DATA_FOLDER/known_classes/test] \\\n",
    "                                       gen_trt_engine.tensorrt.calibration.cal_cache_file=$RESULTS_DIR/gen_trt_engine/cal_resnet101_model.int8.bin\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Generated tensorrt engines and calibration files:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/gen_trt_engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate with generated TensorRT engine\n",
    "!tao deploy ml_recog evaluate -e $SPECS_DIR/evaluate.yaml \\\n",
    "                                 evaluate.trt_engine=$RESULTS_DIR/gen_trt_engine/resnet101_model.engine \\\n",
    "                                 results_dir=$RESULTS_DIR \\\n",
    "                                 dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "                                 dataset.val_dataset.query=$DATA_DIR/$DATA_FOLDER/known_classes/test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference with generated TensorRT engine\n",
    "!tao deploy ml_recog inference -e $SPECS_DIR/infer.yaml \\\n",
    "                                  inference.trt_engine=$RESULTS_DIR/gen_trt_engine/resnet101_model.engine \\\n",
    "                                  results_dir=$RESULTS_DIR \\\n",
    "                                  dataset.val_dataset.reference=$DATA_DIR/$DATA_FOLDER/known_classes/reference \\\n",
    "                                  inference.input_path=$DATA_DIR/$DATA_FOLDER/known_classes/test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('TensorRT Inference results:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/trt_inference"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
